import matplotlib
import os
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import savgol_filter

def readrate(fname):
    infile = open(fname,'r')
    temp = infile.readline()

    while not 'NATOM' in temp:
        temp = infile.readline()
    temp = infile.readline().split() # vals
    natom, nmol, nreact, nfix, nvary = int(temp[0]), int(temp[1]), int(temp[2]), int(temp[4]), int(temp[6])
    temp = infile.readline() # header
    temp = infile.readline().split() # vals 
    nz = int(temp[2])

    # skip all lines that come before 'MIXING RATIO'
    temp = infile.readline()
    while 'REACTION RATES' not in temp:
        temp = infile.readline()

    # loop over groups of species                   
    mix = {}
    nz = 60
    alt = np.zeros(nz)

    blank = infile.readline()
    header = infile.readline()
    header = header.replace('RATE(','')
    header = header.split()
    alt = np.zeros(nz)
    while len(header)>1:
        for ix in range(1,len(header)):
            mix[header[ix]]= np.zeros(nz)
        for iz in range(0,nz):
            temp = infile.readline().split()
            for ix in range(1,len(header)):
                mix[header[ix]][iz] = float(temp[ix+1])
            alt[iz] = float(temp[1])
        blank = infile.readline()
        col = infile.readline()
        blank = infile.readline()
        header = infile.readline()
        header = header.replace('RATE(','')
        header = header.split()
    mix['z'] = alt
    return mix

mypath = ['outputs_1x']#,'outputs_3x','outputs_10x']
dat = {}
for i in mypath:
    dat[i] = {}
    for ifile in os.listdir('./withsep/'+i+'/'):
        if not ifile.endswith('mini'):
            continue
        dat[i][ifile] = readrate('./withsep/'+i+'/'+ifile)

matplotlib.rc('font', family='serif')
matplotlib.rc('xtick', labelsize = 18)
matplotlib.rc('ytick', labelsize = 18)
matplotlib.rc('axes', linewidth=2)
fontname = 'serif'
fontsize=30
lw = 3.3


pre = ['7mbar','50mbar','200mbar','500mbar','1bar']
#species = ['NO','HNO2','HNO3','HO2NO2','NO2','N2O','N2O5']
#species = ['NO','NO2','N2O']
#species = ['23)','24)','25)','26)']#,'101)','113)','114)','137)']
#species = ['118)','121)','123)','131)']
#species = ['117)','139)','140)','142)']
species = ['101)','113)','102)','137)','114)','135)','115)']

#species = ['HNO2','HNO3','HO2NO2','N2O5']
clr = ['royalblue','limegreen','darkorange','brown','mediumvioletred','grey','purple']
#clr = ['red','blue','green','purple','brown','orange','cyan','grey']
lns = ['-','--',':']
fig,ax=plt.subplots(ncols=5,nrows=1,figsize=(24,8))

for i in mypath:
    if  '1x' in i:
        ictr = 0
    #elif '3x' in i:
    #    ictr = 1
    #elif '10x' in i:
    #    ictr = 2
    else:
        print('ice case not found')
    for ifile in list(dat[i].keys()):
        kctr = 0

        if pre[0] in ifile:
            jctr = 0
        elif pre[1] in ifile:
            jctr = 1
        elif pre[2] in ifile:
            jctr = 2
        elif pre[3] in ifile:
            jctr = 3
        elif pre[4] in ifile:
            jctr = 4
        else:
            print('pressure not found, ifile')

        if jctr < 4:
            for isp in list(species):
                ax[jctr].semilogx(dat[i][ifile][isp],dat[i][ifile]['z'],color=clr[kctr],linestyle=lns[ictr],linewidth=2)
                kctr += 1
        else:
            for isp in list(species):
                try:
                    yhat = savgol_filter(np.log10(dat[i][ifile][isp]), 13, 3)
                    yhat2 = savgol_filter(yhat, 15, 3)
                except:
                    yhat2 = np.log10(dat[i][ifile][isp])
                ax[jctr].semilogx(10**yhat2,dat[i][ifile]['z'],color=clr[kctr],linestyle=lns[ictr],linewidth=2)
                kctr += 1

        ax[jctr].set_xlim([1e-6,1e3])
        jctr += 1
    ictr += 1

print(list(dat['outputs_1x'].keys()))
#print(list(dat['outputs_3x'].keys()))
#print(list(dat['outputs_10x'].keys()))
print(list(dat.keys()))

ax[0].legend()
plt.savefig('prodloss.png',dpi=300)

